#中文歌词生成模型
from transformers import GPT2LMHeadModel, BertTokenizer, TextGenerationPipeline

cache_dir = "../../my_model_cache/gpt2-chinese"
# tokenizer = BertTokenizer.from_pretrained("uer/gpt2-chinese-lyric", cache_dir=cache_dir)
# model = GPT2LMHeadModel.from_pretrained("uer/gpt2-chinese-lyric", cache_dir=cache_dir)
tokenizer = BertTokenizer.from_pretrained(
    cache_dir + r"/models--uer--gpt2-chinese-lyric/snapshots/4a42fd76daab07d9d7ff95c816160cfb7c21684f")
model = GPT2LMHeadModel.from_pretrained(
    cache_dir + r"/models--uer--gpt2-chinese-lyric/snapshots/4a42fd76daab07d9d7ff95c816160cfb7c21684f")
#创建模型推理对象
text_generator = TextGenerationPipeline(model, tokenizer)
out = text_generator("在下雨的天，你走在前面", max_length=100, do_sample=True)
print(out)
